import numpy as np
import os
import json
import h5py
import hdf5plugin
import argparse 
from tqdm import tqdm


pku_scenes = ['normal']
pku_splits = ['train', 'val']


event_widht = 1280
event_height = 720


def get_number_in_str(file_name):
    numeric_part = ''.join(char for char in file_name if char.isnumeric())
    return int(numeric_part)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--txt_input_path', type=str, default='E:/Datasets/lastdata/annotations', help='josn label path')
    parser.add_argument('--h5_input_path', type=str, default='E:/Datasets/lastdata/H5', help='raw event path')
    args = parser.parse_args()


    delete_num = 0
    for mode in pku_splits:
        each_mode_josn_dict_path = args.txt_input_path + '/' + mode
        each_mode_h5_dict_path = args.h5_input_path + '/' + mode
        for scene in pku_scenes:
            each_mode_scene_josn_dict_path = each_mode_josn_dict_path + '/' + scene
            each_mode_scene_h5_dict_path = each_mode_h5_dict_path + '/' + scene
            print('prcessing the ' + scene + ' scene in ' + mode)
            for filename in tqdm(os.listdir(each_mode_scene_josn_dict_path)):
                josn_file_path = os.path.join(each_mode_scene_josn_dict_path, filename)
                josn_file_path_ = josn_file_path + '/annotations'
                if os.path.exists(josn_file_path_):
                    josn_file_path = josn_file_path_
                else:
                    josn_file_path = josn_file_path + '/annotation'
                h5_file_path = each_mode_scene_h5_dict_path + '/' + filename + '.raw.h5'
                init_time = each_mode_scene_h5_dict_path + '/init_time/' + filename + '.raw.init_time.h5'
                save_path = each_mode_scene_h5_dict_path + '/' + filename
                file_list = os.listdir(josn_file_path)
                # npy_file = np.load(npy_path)
                if not os.path.exists(h5_file_path):
                    continue
                print(os.path.exists(h5_file_path))
                h5_file = h5py.File(h5_file_path)
                init_time = h5py.File(init_time)


                init_time = init_time['init_time']['time'][()]


                annotations_list = []
                for josn_file_name in file_list:


                    if josn_file_name == 'classes.txt':
                        continue
                    ann_data_list = []
                    file_path = os.path.join(josn_file_path, josn_file_name)


                    with open(file_path, 'r') as file:
                        lines = file.readlines()
                        for line in lines:
                            print(file_path)
                            ann_data = [float(num_str) for num_str in line.strip().split()]
                            ann_data_list.append(ann_data)


                        ann_data_list = np.array(ann_data_list)


                        bbox_num = len(ann_data_list)


                        timestamp = josn_file_name.split(".")[0]
                        timestamp = int(timestamp.split("_")[-1])


                        timestamp = timestamp - init_time


                        for bbox_index in range(bbox_num):
                            x_tl = (ann_data_list[bbox_index][1] - 0.5 * ann_data_list[bbox_index][3]) * event_widht
                            y_tl = (ann_data_list[bbox_index][2] - 0.5 * ann_data_list[bbox_index][4]) * event_height


                            width = ann_data_list[bbox_index][3] * event_widht
                            height = ann_data_list[bbox_index][4] * event_height


                            if width == 0 or height == 0:
                                print(file_path, ' has the width or height = 0')


                            class_id = ann_data_list[bbox_index][0]
                            annotation = (timestamp, x_tl, y_tl, width, height, class_id)
                            annotations_list.append(annotation)


                dtype = [('t', np.uint64), ('x', np.float32), ('y', np.float32), ('w', np.float32), ('h', np.float32), ('class_id', np.uint8)]


                annotations_list_numpy = np.array(annotations_list, dtype=dtype)
                sorted_data_array = np.sort(annotations_list_numpy, order='t')
                if len(annotations_list_numpy) == 0:
                    deleted_file = save_path + '.raw.h5'
                    try:
                        if os.path.exists(deleted_file):
                            os.remove(deleted_file)
                    except PermissionError:
                        print(f"Error: The file {deleted_file} is being used by another program. Please close the program using it and try again.")
                else:
                    save_path = save_path + '_bbox'
                    np.save(save_path, sorted_data_array)
    print('delete {} wrong labels'.format(delete_num))